import argparse
import json
import random
import xml.etree.ElementTree as ET
from pathlib import Path

import fire


def read_xml_label(label_file, image_dir, label_name_set):
    tree = ET.parse(label_file)
    root = tree.getroot()

    label = {}
    label["image"] = str(image_dir / root.find("filename").text)
    targets = []
    for obj in root.iter("object"):
        clss = obj.find("name").text
        bbox = obj.find("bndbox")
        xmin = float(bbox.find("xmin").text)
        ymin = float(bbox.find("ymin").text)
        xmax = float(bbox.find("xmax").text)
        ymax = float(bbox.find("ymax").text)
        targets.append({"type": clss, "bbox": [xmin, ymin, xmax, ymax]})
        label_name_set.add(clss)
    label["label"] = targets

    assert len(targets) > 0
    return label


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_root", type=str, default=None)
    parser.add_argument("--data_split", type=str, default="VOC2007", nargs="+")
    parser.add_argument("--output", type=str, default=None)
    parser.add_argument("--train_ratio", type=float, default=0.85)
    return parser


def voc2json(data_root, data_splits, output_root, train_ratio=0.85):
    labels = []
    label_name_set = set()

    for split in data_splits:
        image_file_dir = Path(data_root) / split / "JPEGImages"
        label_file_dir = Path(data_root) / split / "Annotations"
        file_list = Path(label_file_dir).glob("*.xml")
        for x in file_list:
            labels.append(read_xml_label(x, image_file_dir, label_name_set))

    label_name_dict = {name: idx for (idx, name) in enumerate(label_name_set)}

    for sample in labels:
        for target in sample["label"]:
            target["type"] = label_name_dict[target["type"]]

    num_train = int(train_ratio * len(labels))
    random.shuffle(labels)
    labels_train = labels[:num_train]
    labels_val = labels[num_train:]

    (output_path := Path(output_root)).mkdir(parents=True, exist_ok=True)
    with (output_path / "voc_label_names.json").open(mode="w") as f:
        json.dump(label_name_dict, f, indent=2)
    with (output_path / "voc_samples_total.json").open(mode="w") as f:
        json.dump(labels, f)
    with (output_path / "voc_samples_train.json").open(mode="w") as f:
        json.dump(labels_train, f)
    with (output_path / "voc_samples_val.json").open(mode="w") as f:
        json.dump(labels_val, f)

    print(f"Total: {len(labels)}, train: {len(labels_train)}, val: {len(labels_val)}")
    return


if __name__ == "__main__":
    fire.Fire(voc2json)
